# -*- coding: utf-8 -*-
"""
Created on Tue Oct 24 09:44:25 2017
修改xml中错误的节点,检查object个数
@author: aia_xml
"""

import xml.etree.ElementTree as ET
import os
import shutil

'''计算文件夹下文件个数'''


def fileCountIn(dir):
    return sum([len(files) for root, dirs, files in os.walk(dir)])


def change_xml_node(xml_file):
    tree = ET.parse(xml_file)
    root = tree.getroot()
    #    for img_name in root.iter('filename'):
    #        img_name.text = img_name.text+'.jpg'

    name_text = 'person'
    obj_num = 0
    chged = 0
    for obj in root.iter('object'):
        if obj.find('name').text == 'object':
            obj.find('name').text = name_text
            chged = 1
        obj_num = obj_num + 1
    if chged == 1:
        print(xml_file, 'has been changed')
    if obj_num == 0:
        print(xml_file, 'has no object')
    tree.write(xml_file)


def main():
    xml_dir = '/home/ccj/dataset/vehicle_person/COCO/voc/Annotations'
    save_dir = '/home/ccj/dataset/test-sets/pedestrian'
    save_annotation_dir = os.path.join(save_dir, "Annotations")
    save_image_dir = os.path.join(save_dir, "JPEGImages")
    save_yolov3_labels = os.path.join(save_dir, "labels")
    if not os.path.exists(save_annotation_dir):
        os.makedirs(save_annotation_dir)

    if not os.path.exists(save_image_dir):
        os.makedirs(save_image_dir)

    xml_Count = fileCountIn(xml_dir)
    print('xml Count in dir:', xml_Count)
    chg_num = 0
    copied_img = 0
    object_dicts = {}
    for root, dirs, files in os.walk(xml_dir):
        for file in files:
            xml_file = os.path.join(xml_dir, file)
            # change_xml_node(xml_file)
            # objects = find_objects(xml_file)

            tree = ET.parse(xml_file)
            root = tree.getroot()
            for obj in root.iter('object'):
                if obj.find('name').text == 'person':
                    xml_name = os.path.split(xml_file)[1]
                    image_name = xml_name.split('.')[0] + '.jpg'
                    image_path = os.path.join(os.path.split(xml_file)[0].replace('Annotations', 'JPEGImages'),
                                              image_name)
                    dst_xml_full_path = os.path.join(save_annotation_dir, xml_name)
                    dst_img_full_path = os.path.join(save_image_dir, image_name)
                    shutil.copyfile(xml_file, dst_xml_full_path)
                    shutil.copyfile(image_path, dst_img_full_path)
                    copied_img += 1
                    break


            chg_num = chg_num + 1
            if chg_num % 1000 == 0:
                print('completed ', chg_num * 100 / xml_Count, '% ...')
    print('completed ', chg_num * 100 / xml_Count, '% ...')
    print('total:', chg_num)
    print('copied image: ', copied_img)


if __name__ == '__main__':
    main()
